package com.crois.barrier.web.filter;


import com.crois.barrier.common.exception.BarrierExceptionEnum;
import com.crois.barrier.web.filter.support.BodyInserterContext;
import com.crois.barrier.web.filter.support.CachedBodyOutputMessage;
import com.crois.barrier.web.resp.BarrierResponseUtils;
import com.crois.barrier.web.resp.BarrierResponseWrapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;

/**
 * @author Hou Ze Yu
 * @description 文件大小过滤器
 * @date 2020/8/19
 */
public class FileSizeFilter implements WebFilter {

    @Value("${door.file.size:10}")
    private int maxSize;

    private static final long BYTES_PER_MB = 1024 * 1024;

    private final List<HttpMessageReader<?>> messageReaders;

    public FileSizeFilter() {
        this.messageReaders = HandlerStrategies.withDefaults().messageReaders();
    }

    @Override
    public Mono<Void> filter(ServerWebExchange serverWebExchange, WebFilterChain webFilterChain) {
        MediaType mediaType = serverWebExchange.getRequest().getHeaders().getContentType();
        if (MediaType.MULTIPART_FORM_DATA.isCompatibleWith(mediaType)){
            ServerRequest serverRequest = ServerRequest.create(serverWebExchange,
                    messageReaders);
            return serverRequest.bodyToMono(DataBuffer.class)
                    .flatMap(size -> {
                        if (size.capacity() > BYTES_PER_MB*maxSize){
                            ServerHttpResponse response = serverWebExchange.getResponse();
                            response.setStatusCode(HttpStatus.BAD_REQUEST);
                            return BarrierResponseUtils.error(serverWebExchange, BarrierResponseWrapper.fail(BarrierExceptionEnum.BARRIER_GATEWAY_FILE_MAX_SIZE));
                        }
                        BodyInserter bodyInserter = BodyInserters.fromPublisher(Mono.just(size), DataBuffer.class);
                        HttpHeaders headers = new HttpHeaders();
                        headers.putAll(serverWebExchange.getRequest().getHeaders());
                        headers.remove(HttpHeaders.CONTENT_LENGTH);
                        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(
                                serverWebExchange, headers);
                        return bodyInserter.insert(outputMessage, new BodyInserterContext())
                                .then(Mono.defer(() -> {
                                    ServerHttpRequest decorator = decorator(serverWebExchange, outputMessage);
                                    return webFilterChain.filter(serverWebExchange.mutate().request(decorator).build());
                                }));
            });
        }
        return webFilterChain.filter(serverWebExchange);
    }

    private ServerHttpRequestDecorator decorator(ServerWebExchange serverWebExchange, CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(serverWebExchange.getRequest()){
            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }
}
